# 华为ModelArts体验太差，无法正常使用

import os
import moxing as mox
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', type=int, default=128, help='batch size')
parser.add_argument('--data_url', type=str, default='obs://***/MLPR/data/', help='data url')
parser.add_argument('--train_url', type=str, default='obs://***/MLPR/model/', help='train url')
parser.add_argument('--init_method', type=str, help='make argparser happy')
parser.add_argument('--num_gpus', type=str, help='make argparser happy')
arg = parser.parse_args()

# Copy Train Data
mox.file.copy_parallel(arg.data_url, '/cache/data/')

# Deps
os.system('pip install tensorboard')
os.system('pip install --upgrade pandas')

# Prepare Data
os.system('unzip /cache/data/RAF-DB/basic/Image/aligned.zip -d /cache/data/RAF-DB/basic/Image/ > /dev/null')
os.system('unzip /cache/data/SFEW/Train/Train_Aligned_Faces.zip -d /cache/data/SFEW/Train/ > /dev/null')
os.system('unzip /cache/data/SFEW/Val/Val_Aligned_Faces_new.zip -d /cache/data/SFEW/Val/ > /dev/null')

# Path
current_path = os.path.dirname(os.path.realpath(__file__))
model_path = os.path.join(current_path, 'model')
os.system('mkdir -p {0}/model && mkdir -p {0}/runs'.format(current_path))

# Run
start = 'python {0}/main.py '.format(current_path)
start += '--rafdb_base="{0}" '.format('/cache/data/RAF-DB/basic/')
start += '--sfew_base="{0}" '.format('/cache/data/SFEW/')
start += '--checkpoint_dir="{0}" '.format(model_path)
start += '--log_dir="{0}" '.format(model_path)
start += '--batch_size={0} '.format(arg.batch_size)
start += '--fuck_huawei={0} '.format(current_path + '/resnet')
start += '--drop_rate=0.4'
os.system(start)

# Copy Back Model
mox.file.copy_parallel(model_path, arg.train_url)
